{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IafxybMjKfBO"
   },
   "source": [
    "##### Copyright 2020 Google"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "pc1aHcGvKmHe"
   },
   "outputs": [],
   "source": [
    "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3707a33c52de"
   },
   "source": [
    "# Quickstart\n",
    "\n",
    "This code tutorial shows how to estimate a 1-RDM and perform variational optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FQEYY3gnK51d"
   },
   "source": [
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://www.example.org/cirq/experiments/hfvqe/quickstart\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on QuantumLib</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/quantumlib/ReCirq/blob/master/docs/hfvqe/quickstart.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/quantumlib/ReCirq/blob/master/docs/hfvqe/quickstart.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://storage.googleapis.com/tensorflow_docs/ReCirq/docs/hfvqe/quickstart.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8625de3d3c0d"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import recirq\n",
    "except ImportError:\n",
    "    !pip install --quiet git+https://github.com/quantumlib/ReCirq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1fac10877d62"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cirq\n",
    "\n",
    "from recirq.hfvqe.gradient_hf import rhf_func_generator\n",
    "from recirq.hfvqe.opdm_functionals import OpdmFunctional\n",
    "from recirq.hfvqe.analysis import (\n",
    "    compute_opdm, mcweeny_purification,\n",
    "    resample_opdm, fidelity_witness,\n",
    "    fidelity)\n",
    "from recirq.hfvqe.third_party.higham import fixed_trace_positive_projection\n",
    "from recirq.hfvqe.molecular_example import make_h6_1_3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "09752498051c"
   },
   "source": [
    "## Set up the experiment\n",
    "\n",
    "Generate the input files, set up quantum resources, and set up the OpdmFunctional to make measurements. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1b750bbaf408"
   },
   "outputs": [],
   "source": [
    "rhf_objective, molecule, parameters, obi, tbi = make_h6_1_3()\n",
    "ansatz, energy, gradient = rhf_func_generator(rhf_objective)\n",
    "\n",
    "# settings for quantum resources\n",
    "qubits = [cirq.GridQubit(0, x) for x in range(molecule.n_orbitals)]\n",
    "sampler = cirq.Simulator(dtype=np.complex128)  # this can be a QuantumEngine\n",
    "\n",
    "# OpdmFunctional contains an interface for running experiments\n",
    "opdm_func = OpdmFunctional(qubits=qubits,\n",
    "                           sampler=sampler,\n",
    "                           constant=molecule.nuclear_repulsion,\n",
    "                           one_body_integrals=obi,\n",
    "                           two_body_integrals=tbi,\n",
    "                           # only simulate spin-up electrons:\n",
    "                           num_electrons=molecule.n_electrons // 2,\n",
    "                           clean_xxyy=True,\n",
    "                           purification=True\n",
    "                           )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "61f9277513fe"
   },
   "source": [
    "The displayed text is the output of the gradient based restricted Hartree-Fock.  We define the gradient in `rhf_objective` and use the conjugate-gradient optimizer to optimize the basis rotation parameters.  This is equivalent to doing Hartree-Fock theory from the canonical transformation perspective."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "61f9277513fe"
   },
   "source": [
    "## Estimate Quantities\n",
    "\n",
    "Next, we will do the following:\n",
    "\n",
    "1. Do measurements for a given set of parameters\n",
    "\n",
    "2. Compute 1-RDM, variances, and purification\n",
    "\n",
    "3. Compute energy, fidelities, and errorbars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8e91bce6cc7c"
   },
   "outputs": [],
   "source": [
    "# 1.\n",
    "# default to 250_000 shots for each circuit.\n",
    "# 7 circuits total, printed for your viewing pleasure\n",
    "# return value is a dictionary with circuit results for each permutation\n",
    "measurement_data = opdm_func.calculate_data(parameters)\n",
    "\n",
    "# 2.\n",
    "opdm, var_dict = compute_opdm(measurement_data, return_variance=True)\n",
    "opdm_pure = mcweeny_purification(opdm)\n",
    "\n",
    "# 3.\n",
    "raw_energies = []\n",
    "raw_fidelity_witness = []\n",
    "purified_eneriges = []\n",
    "purified_fidelity_witness = []\n",
    "purified_fidelity = []\n",
    "true_unitary = ansatz(parameters)\n",
    "nocc = molecule.n_electrons // 2\n",
    "nvirt = molecule.n_orbitals - nocc\n",
    "initial_fock_state = [1] * nocc + [0] * nvirt\n",
    "\n",
    "# 1000 repetitions of the measurement\n",
    "for _ in range(1000):  \n",
    "    new_opdm = resample_opdm(opdm, var_dict)\n",
    "    raw_energies.append(opdm_func.energy_from_opdm(new_opdm))\n",
    "    raw_fidelity_witness.append(\n",
    "        fidelity_witness(target_unitary=true_unitary,\n",
    "                         omega=initial_fock_state,\n",
    "                         measured_opdm=new_opdm)\n",
    "    )\n",
    "    # fix positivity and trace of sampled 1-RDM if strictly outside\n",
    "    # feasible set\n",
    "    w, v = np.linalg.eigh(new_opdm)\n",
    "    if len(np.where(w < 0)[0]) > 0:\n",
    "        new_opdm = fixed_trace_positive_projection(new_opdm, nocc)\n",
    "\n",
    "    new_opdm_pure = mcweeny_purification(new_opdm)\n",
    "    purified_eneriges.append(opdm_func.energy_from_opdm(new_opdm_pure))\n",
    "    purified_fidelity_witness.append(\n",
    "        fidelity_witness(target_unitary=true_unitary,\n",
    "                         omega=initial_fock_state,\n",
    "                         measured_opdm=new_opdm_pure)\n",
    "    )\n",
    "    purified_fidelity.append(\n",
    "        fidelity(target_unitary=true_unitary,\n",
    "                 measured_opdm=new_opdm_pure)\n",
    "    )\n",
    "print(\"Canonical Hartree-Fock energy \", molecule.hf_energy)\n",
    "print(\"True energy \", energy(parameters))\n",
    "print(\"Raw energy \", opdm_func.energy_from_opdm(opdm),\n",
    "      \"+- \", np.std(raw_energies))\n",
    "print(\"Raw fidelity witness \", np.mean(raw_fidelity_witness).real,\n",
    "      \"+- \", np.std(raw_fidelity_witness))\n",
    "print(\"purified energy \", opdm_func.energy_from_opdm(opdm_pure),\n",
    "      \"+- \", np.std(purified_eneriges))\n",
    "print(\"Purified fidelity witness \", np.mean(purified_fidelity_witness).real,\n",
    "      \"+- \", np.std(purified_fidelity_witness))\n",
    "print(\"Purified fidelity \", np.mean(purified_fidelity).real,\n",
    "      \"+- \", np.std(purified_fidelity))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "669a320cb246"
   },
   "source": [
    "This prints out the various energies estimated from the 1-RDM along with error bars.  Generated from resampling the 1-RDM based on the estimated covariance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7306afa48121"
   },
   "source": [
    "## Optimization\n",
    "\n",
    "We use the sampling functionality to variationally relax the parameters of\n",
    "my ansatz such that the energy is decreased.\n",
    "\n",
    "For this we will need the augmented Hessian optimizer\n",
    "\n",
    "The optimizerer code we have takes:\n",
    "rhf_objective object, initial parameters,\n",
    "a function that takes a n x n unitary and returns an opdm\n",
    "maximum iterations,\n",
    "hassian_update which indicates how much of the hessian to use\n",
    "rtol which is the gradient stopping condition.\n",
    "\n",
    "A natural thing that we will want to save is the variance dictionary of\n",
    "the non-purified 1-RDM.  This is accomplished by wrapping the 1-RDM\n",
    "estimation code in another object that keeps track of the variance \n",
    "dictionaries. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9d40a667216a"
   },
   "outputs": [],
   "source": [
    "from recirq.hfvqe.mfopt import moving_frame_augmented_hessian_optimizer\n",
    "from recirq.hfvqe.opdm_functionals import RDMGenerator\n",
    "\n",
    "rdm_generator = RDMGenerator(opdm_func, purification=True)\n",
    "opdm_generator = rdm_generator.opdm_generator\n",
    "\n",
    "result = moving_frame_augmented_hessian_optimizer(\n",
    "    rhf_objective=rhf_objective,\n",
    "    initial_parameters=parameters + 1.0E-1,\n",
    "    opdm_aa_measurement_func=opdm_generator,\n",
    "    verbose=True, delta=0.03,\n",
    "    max_iter=20,\n",
    "    hessian_update='diagonal',\n",
    "    rtol=0.50E-2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bfe90221028b"
   },
   "source": [
    "Each interation prints out a variety of information that the user might find useful.  Watching energies go down is known to be one of the best forms of entertainment during a shelter-in-place order.\n",
    "\n",
    "After the optimization we can print the energy as a function of iteration number to see close the energy gets to the true minium."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4119e4ea8417"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.semilogy(range(len(result.func_vals)),\n",
    "             np.abs(np.array(result.func_vals) - energy(parameters)),\n",
    "             'C0o-')\n",
    "plt.xlabel(\"Optimization Iterations\",  fontsize=18)\n",
    "plt.ylabel(r\"$|E  - E^{*}|$\", fontsize=18)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "quickstart.ipynb",
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
